import sys
import time

import numpy as np
import random
import pandas as pd
import networkx as nx
from tqdm import tqdm
import argparse
from src.RKGE import train
from keras.models import load_model

import get_string
from src.load_base import load_ratings, data_split, get_records, get_rec, getentity_list, get_trueid


def get_file_line_count(file_path):
    line_count = 0
    with open('entity_list', 'r') as file:
        for line in file:
            line_count += 1
    return line_count

def construct_kg(data_dir, train_records, user_set):

    kg = nx.DiGraph()
    triple_np = pd.read_csv(data_dir + 'kg.txt', delimiter='\t', header=None).values
    # 获得起始点构成的集合
    head_set = set(triple_np[:, 0])
    tail_set = set(triple_np[:, 1])
    entity_list = list(head_set.union(tail_set)) + list(user_set)





    for triple in triple_np:
        head = entity_list.index(triple[0])
        tail = entity_list.index(triple[1])
        relation = triple[2]

        kg.add_edge(head, tail, relation=relation)


    n_relation = len(set(triple_np[:, 2]))

    for user in train_records:
        for item in train_records[user]:
            kg.add_edge(entity_list.index(user), entity_list.index(item), relation=n_relation)
            kg.add_edge(entity_list.index(item), entity_list.index(user), relation=n_relation + 1)
    return kg, entity_list

def get_relation_dict(kg):

    relation_dict = dict()
    for edge in tqdm(kg.edges):

        relation_dict[(edge[0], edge[1])] = kg[edge[0]][edge[1]]['relation']

    return relation_dict


def get_paths(args,zfc):
    np.random.seed(555)
    # np.random.seed(int(time.time()))
    data_dir = './data/' + args.dataset + '/'
    ratings_np = load_ratings(data_dir)
    entity_num=getentity_list(data_dir)
    train_set, eval_set, test_set = data_split(ratings_np, args.ratio,entity_num)
    item_set = set(ratings_np[:, 1])
    user_set = set(ratings_np[:, 0])
    train_records = get_records(train_set)
    # 获取的都是正样品
    eval_records = get_records(eval_set)
    test_records = get_records(test_set)
    rec = get_rec(train_records, eval_records, test_records, item_set)
    # ??????????????????????
    kg, entity_list = construct_kg(data_dir, train_records, user_set)
    relation_dict = get_relation_dict(kg)

    path_dict = dict()
    new_train_set = []
    # 这段代码是使用networkx库中的all_simple_paths函数找出从user节点到item节点的所有简单路径，并将路径长度为args.path_len + 1的路径筛选出来，最终将这些路径存储到列表paths中。
    #
    # 在具体实现中，all_simple_paths函数的第一个参数kg表示一个有向图对象，第二个参数user表示起始节点，第三个参数item表示终止节点，cutoff参数
    # 表示最大路径长度限制，如果不指定，默认为None，表示不限制最大路径长度。列表推导式的过滤条件是使用if语句实现的，如果一个路径的长度为
    # args.path_len + 1，即从user到item需要经过args.path_len个中间节点，则将该路径添加到列表paths中。
    for pair in tqdm(train_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_train_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_eval_set = []

    for pair in tqdm(eval_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_eval_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if
                 (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_test_set = []

    for pair in tqdm(test_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_test_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if
                 (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_rec = dict()
    for user in tqdm(rec):
        new_rec[entity_list.index(user)] = [entity_list.index(i) for i in rec[user]]
        new_user = entity_list.index(user)
        for item in new_rec[new_user]:

            paths = [path for path in list(nx.all_simple_paths(kg, new_user, item, cutoff=args.path_len)) if (len(path) == args.path_len + 1)]

            if len(paths) > 50:
                indices = np.random.choice(len(paths), 50, replace=False)
                paths = [paths[i] for i in indices]

            path_dict[(new_user, item)] = paths

    # true_id=get_trueid(data_dir,zfc)
    # data = np.zeros((165, 3))  # 创建一个1975行、3列的全0数组
    #
    # for i in range(165):
    #     data[i, 0] = true_id
    #     data[i, 1] = i
    #     data[i, 2] = random.randint(0, 1)

    np.save(data_dir+str(args.ratio)+'_'+str(args.path_len)+'_path_dict.npy', path_dict)
    np.save(data_dir+str(args.ratio)+'_relation_dict.npy', relation_dict)
    np.save(data_dir+'_entity_list.npy', entity_list)
    np.save(data_dir + str(args.ratio)+'_train_set.npy', new_train_set)
    np.save(data_dir + str(args.ratio) + '_eval_set.npy', new_eval_set)
    np.save(data_dir + str(args.ratio)+'_test_set.npy', new_test_set)
    np.save(data_dir + str(args.ratio)+'_rec.npy', new_rec)
    # # 添加代码
    # np.save(data_dir + str(args.ratio) + '_test.npy', data)
    # # 添加代码

def get_paths1(args,zfc):
    # data_dir = './RKGE-master/data/' + args.dataset + '/'
    data_dir = './data/' + args.dataset + '/'
    # 添加代码
    data = np.zeros((164, 3))  # 创建一个1975行、3列的全0数组

    for i in range(164):
        data[i, 0] = zfc
        data[i, 1] = i
        data[i, 2] = random.randint(0, 1)


    # 添加代码
    np.save(str(data_dir) + str(args.ratio) + '_test.npy', data)
    # 添加代码

def recommend_attractions(district, purpose, type, season, age, transportation):
    # district=str(district)
    # purpose = str(purpose)
    # type = str(type)
    # season = str(season)
    # age=str(age)
    # transportation = str(transportation)
    #
    # params = district + purpose + type + season + age + transportation
    params=district*2880+purpose*720+type*144+season*24+age*4+transportation;
    return params


class Configuration:
    def __init__(self):
        self.dataset = 'zhishitupu'
        self.ratio = 0.2
        self.lr = 5e-3
        self.l2 = 1e-4
        self.batch_size = 1024
        self.device = 'cuda:0'
        self.dim = 20
        self.p = 5

# 创建一个配置对象
config = Configuration()

# # 通过对象访问参数
# print("Dataset:", config.dataset)
# print("Ratio:", config.ratio)
# print("Learning Rate:", config.lr)
# print("L2:", config.l2)
# print("Batch Size:", config.batch_size)
# print("Device:", config.device)
# print("Dimension:", config.dim)
# print("Number of Paths:", config.p)

if __name__ == "__main__":
    district=0
    purpose=0
    type=0
    season=0
    age=0
    transportation=0

    # district = sys.argv[1]
    # purpose = sys.argv[2]
    # type=sys.argv[3]
    # season=sys.argv[4]
    # age=sys.argv[5]
    # transportation=sys.argv[6]
    zfc = recommend_attractions(district, purpose, type, season, age, transportation)
    get_paths1(config, zfc)
    train(config, True)


